package polygon.util;

import org.apache.log4j.Logger;

import java.io.*;
import java.util.Arrays;
import java.util.Enumeration;
import java.util.Random;
import java.util.zip.*;

/**
 * @author Mike Mirzayanov
 */
public class ZipUtil {
    private static final Logger logger = Logger.getLogger(ZipUtil.class);
    private static final long MAX_UNCOMPRESS_FILE_SIZE = 10 * 1024 * 1024;
    private static final long MAX_UNCOMPRESS_ENTRY_COUNT = 10000;

    public static byte[] decompress(byte[] bytes) throws DataFormatException {
        Inflater decompresser = new Inflater();
        decompresser.setInput(bytes);

        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        byte[] buffer = new byte[bytes.length * 2 + 128];

        while (!decompresser.finished()) {
            outputStream.write(buffer, 0, decompresser.inflate(buffer));
        }

        decompresser.end();
        return outputStream.toByteArray();
    }

    public static byte[] compress(byte[] bytes) {
        Deflater compressor = new Deflater();
        compressor.setInput(bytes);
        compressor.finish();

        ByteArrayOutputStream outputStream = new ByteArrayOutputStream(bytes.length);

        byte[] buffer = new byte[bytes.length * 2 + 128];

        while (!compressor.finished()) {
            outputStream.write(buffer, 0, compressor.deflate(buffer));
        }

        try {
            outputStream.close();
        } catch (IOException e) {
            // No operations.
        }

        return outputStream.toByteArray();
    }

    public static void unzip(byte[] bytes, File directory) throws IOException {
        File zipped = File.createTempFile("zip", System.currentTimeMillis() + "");
        FileUtil.writeFile(zipped, bytes);

        try {
            ZipFile zipFile = new ZipFile(zipped);
            Enumeration entries = zipFile.entries();

            int count = 0;

            while (entries.hasMoreElements() && count < MAX_UNCOMPRESS_ENTRY_COUNT) {
                count++;
                ZipEntry entry = (ZipEntry) entries.nextElement();

                if (entry.isDirectory()) {
                    File dir = new File(directory, entry.getName());
                    if (!dir.mkdir()) {
                        throw new IOException("Can't create directory while unzip process.");
                    }
                    continue;
                }

                if (entry.getSize() <= MAX_UNCOMPRESS_FILE_SIZE && entry.getCompressedSize() <= MAX_UNCOMPRESS_FILE_SIZE) {
                    File file = new File(directory, entry.getName());
                    if (!file.getParentFile().isDirectory()) {
                        if (!file.getParentFile().mkdirs()) {
                            throw new IOException("Can't create directory " + file.getParentFile() + ".");
                        }
                    }
                    copyInputStream(zipFile.getInputStream(entry),
                            new BufferedOutputStream(new FileOutputStream(file)));
                } else {
                    long size = Math.max(entry.getSize(), entry.getCompressedSize());
                    throw new IOException("File " + entry.getName() + " too large: " + size + " bytes.");
                }
            }

            zipFile.close();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (!zipped.delete()) {
                logger.error("Can't delete file " + zipped + ".");
            }
        }
    }

    private static void copyInputStream(InputStream in, OutputStream out) throws IOException {
        byte[] buffer = new byte[1024];
        int len;

        while ((len = in.read(buffer)) >= 0) {
            out.write(buffer, 0, len);
        }

        in.close();
        out.close();
    }

    public static void zip(File source, File destination) throws IOException {
        ZipOutputStream zipOutputStream = new ZipOutputStream(new FileOutputStream(destination));
        zipOutputStream.setLevel(9);

        try {
            addDirectory("", source, zipOutputStream);
        } finally {
            zipOutputStream.close();
        }
    }

    private static void addDirectory(String prefix, File source, ZipOutputStream zipOutputStream) throws IOException {
        File[] files = source.listFiles();

        for (File file : files) {
            if (!file.isHidden()) {
                if (file.isDirectory()) {
                    String dir = prefix + file.getName() + "/";
                    zipOutputStream.putNextEntry(new ZipEntry(dir));
                    zipOutputStream.closeEntry();
                    addDirectory(dir, file, zipOutputStream);
                }
            }
        }

        for (File file : files) {
            if (!file.isHidden()) {
                if (file.isFile()) {
                    String path = prefix + file.getName();
                    zipOutputStream.putNextEntry(new ZipEntry(path));
                    zipOutputStream.write(FileUtil.getBytes(file));
                    zipOutputStream.closeEntry();
                }
            }
        }
    }

    public static void main(String[] args) throws DataFormatException {
        Random random = new Random();

        for (int i = 0; i < 10000000; i += 13131) {
            int length = i + 1;
            byte[] bytes = new byte[length];
            for (int j = 0; j < length; j++) {
                bytes[j] = (byte) random.nextInt(256);
            }
            byte[] compressed = compress(bytes);
            byte[] uncompressed = decompress(compressed);
            if (!Arrays.equals(bytes, uncompressed)) {
                throw new RuntimeException();
            }
            System.out.println(i + ": " + bytes.length + " " + compressed.length);
        }
    }
}
